package cn.doitedu.ml.examples

import cn.doitedu.commons.utils.SparkUtil
import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.linalg.Vectors

import scala.collection.mutable

/**
 * @Title: ${file_name}
 * @Package ${package_name}
 * @Description: ${todo}
 * @Author hunter@doitedu.cn
 * @date 2020/8/1615:35     
 */
object BayesDemo {

  def main(args: Array[String]): Unit = {

    val spark = SparkUtil.getSparkSession("朴素贝叶斯算法示例：预测出轨概率")
    import spark.implicits._
    import org.apache.spark.sql.functions._

    // 加载原始样本集
    val sample = spark.read.option("header", "true").csv("portrait/testdata/bayes/sample/sample.csv")
    sample.createTempView("sample")
    // 特征值数字化
    val digitFeatures = spark.sql(
      """
        |
        |select
        |name,
        |case job
        |  when '老师' then 0.0
        |  when '程序员' then 1.0
        |  else 2.0
        |end as job,
        |
        |case income
        |  when '中' then 0.0
        |  when '高' then 1.0
        |  else 2.0
        |end as income,
        |
        |case sex
        |  when '男' then 0.0
        |  when '女' then 1.0
        |  else 2.0
        |end as sex,
        |
        |case age
        |  when '青年' then 0.0
        |  when '老年' then 1.0
        |  else 2.0
        |end as age,
        |
        |case amour
        |  when '出轨' then 0.0
        |  when '没出' then 1.0
        |  else 2.0
        |end as label
        |
        |from sample
        |
        |""".stripMargin)
    digitFeatures.show(50,false)

    // 定义一个将数组转成Vector的函数
    val arr2Vec = udf((arr:mutable.WrappedArray[Double])=>{
      Vectors.dense(arr.toArray)
    })

    // 特征向量化
    val featureVecs = digitFeatures.select('name,'label,arr2Vec(array('job,'income,'age,'sex)) as "vec")

    // 构造算法工具
    val naiveBayes = new NaiveBayes()
        .setFeaturesCol("vec")
        .setLabelCol("label")
        .setSmoothing(1.0)

    // 针对样本集训练模型
    val model = naiveBayes.fit(featureVecs)
    model.save("portrait/testdata/bayes/model")

    // 加载训练好的模型
    val model1 = NaiveBayesModel.load("portrait/testdata/bayes/model")

    // 加载测试数据集
    val test = spark.read.option("header", "true").csv("portrait/testdata/bayes/test")
    test.createTempView("test")

    spark.udf.register("arr2Vec",arr2Vec)
    val featureVecTest = spark.sql(
      """
        |
        |select
        |name,
        |arr2Vec(
        |array(
        |case job
        |  when '老师' then 0.0
        |  when '程序员' then 1.0
        |  else 2.0
        |end,
        |
        |case income
        |  when '中' then 0.0
        |  when '高' then 1.0
        |  else 2.0
        |end,
        |
        |case sex
        |  when '男' then 0.0
        |  when '女' then 1.0
        |  else 2.0
        |end,
        |
        |case age
        |  when '青年' then 0.0
        |  when '老年' then 1.0
        |  else 2.0
        |end
        |)
        |) as vec
        |
        |
        |from test
        |
        |
        |""".stripMargin)

    val res = model1.transform(featureVecTest)
    res.show(100,false)

    spark.close()
  }

}
